import logging
from typing import Any, Dict, List, Optional, Tuple

try:
    from .vector_storage import ElasticStorage, MilvusStorage
except ImportError:
    # 兼容脚本直跑
    from utils.use_file.vector_storage import ElasticStorage, MilvusStorage


logger = logging.getLogger(__name__)


class HybridGuideSearch:
    """
    混合检索门面：聚合关键词检索（Elasticsearch）与向量检索（Milvus），
    提供统一的索引写入与融合排序检索接口。

    用法概览：
    - 初始化：hs = HybridGuideSearch()
    - 准备：hs.setup(recreate=False)  # 可选，仅为确保 ES 索引存在
    - 写入：hs.index_blocks(blocks, embeddings)
      blocks 形如 [{"content": str, "metadata": {...}}]，
      embeddings 与 blocks 对齐，元素为 List[float]
    - 查询：hs.hybrid_search(query, query_embedding, size, ...)
    """
    def __init__(self, es: Optional[ElasticStorage] = None, milvus: Optional[MilvusStorage] = None) -> None:
        self.es = es or ElasticStorage()
        self.milvus = milvus or MilvusStorage()

    def setup(self, recreate: bool = False) -> None:
        """确保 ES 索引存在；当 recreate=True 时重建索引。

        说明：Milvus 的集合在 MilvusStorage 初始化时已自检并创建。
        """
        self.es.ensure_index(recreate=recreate)

    def index_blocks(
        self,
        blocks: List[Dict[str, Any]],
        embeddings: Optional[List[List[float]]] = None,
    ) -> None:
        """批量入库文本块与向量。

        - 先写 ES 文档（便于关键词检索与字段过滤）
        - 再写 Milvus 向量（用于相似度检索）
        """
        # 写入 ES 文档
        self.es.upsert_documents(blocks)

        # 写入 Milvus 向量
        if embeddings is not None:
            items: List[Dict[str, Any]] = []
            for block, vec in zip(blocks, embeddings):
                meta = block.get("metadata", {})
                items.append({
                    "text_block_id": meta.get("text_block_id"),
                    "embedding": vec,
                    "metadata": meta,
                })
            self.milvus.upsert_vectors(items)

    def hybrid_search(
        self,
        query: str,
        query_embedding: Optional[List[float]] = None,
        size: int = 10,
        weight_es: float = 0.5,
        weight_vec: float = 0.5,
        user_role: Optional[str] = None,
        chapter: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """执行融合检索并返回排序后的结果。

        - 先各自查询：ES 文本匹配；Milvus 向量相似度
        - 对两路分数做 0-1 归一化
        - 按 weight_es、weight_vec 加权融合
        - 拉取 ES 正文作为最终输出
        """
        # 关键词检索 (ES)
        es_results = self.es.search(
            query=query,
            size=size,
            user_role=user_role,
            chapter=chapter,
            content_type=content_type,
        )
        es_by_id = {r.get("text_block_id"): r for r in es_results}

        # 向量检索 (Milvus)
        milvus_results: List[Dict[str, Any]] = []
        if query_embedding is not None:
            milvus_results = self.milvus.search(
                embedding=query_embedding,
                top_k=size,
                user_role=user_role,
                chapter=chapter,
                content_type=content_type,
            )

        def normalize(scores: List[float]) -> Dict[Any, float]:
            """对同一路结果的原始分数做区间缩放，避免量纲不一致。"""
            if not scores:
                return {}
            mn = min(scores)
            mx = max(scores)
            if mx - mn < 1e-9:
                return {i: 1.0 for i, _ in enumerate(scores)}
            return {i: (s - mn) / (mx - mn) for i, s in enumerate(scores)}

        es_scores = [r.get("_es_score", 0.0) for r in es_results]
        es_norm = normalize(es_scores)

        mv_scores = [r.get("_vector_score", 0.0) for r in milvus_results]
        mv_norm = normalize(mv_scores)

        mv_by_id = {r.get("text_block_id"): r for r in milvus_results}

        all_ids = set(es_by_id.keys()) | set(mv_by_id.keys())
        merged: List[Tuple[str, float]] = []
        for tid in all_ids:
            es_idx = next((i for i, r in enumerate(es_results) if r.get("text_block_id") == tid), None)
            mv_idx = next((i for i, r in enumerate(milvus_results) if r.get("text_block_id") == tid), None)

            es_score = es_norm.get(es_idx, 0.0) if es_idx is not None else 0.0
            mv_score = mv_norm.get(mv_idx, 0.0) if mv_idx is not None else 0.0
            merged.append((tid, weight_es * es_score + weight_vec * mv_score))

        merged.sort(key=lambda x: x[1], reverse=True)
        top_ids = [tid for tid, _ in merged[:size]]

        id_to_doc = self.es.get_by_ids(top_ids)
        results: List[Dict[str, Any]] = []
        for tid in top_ids:
            doc = id_to_doc.get(tid)
            if not doc:
                mv = mv_by_id.get(tid, {})
                doc = {
                    "text_block_id": tid,
                    "user_role": mv.get("user_role", ""),
                    "chapter": mv.get("chapter", ""),
                    "content_type": mv.get("content_type", ""),
                    "pdf_page": mv.get("pdf_page", 0),
                    "content": "",
                }
            results.append(doc)

        return results


if __name__ == "__main__":
    try:
        from .text_embedding import TextEmbedding
    except ImportError:
        from utils.use_file.text_embedding import TextEmbedding

    logging.basicConfig(level=logging.INFO)

    hs = HybridGuideSearch()
    embedder = TextEmbedding()
    # 固定查询与参数
    user_question = "图片说明 企业注册 操作步骤"
    topk = 5
    user_role = None
    chapter = None
    content_type = None

    # 生成查询向量；若只想看 ES，可不传入向量改用 None
    q_emb = embedder.generate_embeddings(user_question)[0]
    results = hs.hybrid_search(
        query=user_question,
        query_embedding=q_emb,
        size=topk,
        user_role=user_role,
        chapter=chapter,
        content_type=content_type,
    )

    # 单独打印 ES 检索结果
    es_only = hs.es.search(
        query=user_question,
        size=topk,
        user_role=user_role,
        chapter=chapter,
        content_type=content_type,
    )
    logging.info("===== ES 检索结果 =====")
    for i, r in enumerate(es_only, 1):
        logging.info(f"ES Top{i} 分数={r.get('_es_score')} 章节={r.get('chapter')} 角色={r.get('user_role')} 页={r.get('pdf_page')}\n内容: {r.get('content','')[:200]}...")

    # 单独打印 Milvus 检索结果（若向量库已构建）
    mv_only = hs.milvus.search(
        embedding=q_emb,
        top_k=topk,
        user_role=user_role,
        chapter=chapter,
        content_type=content_type,
    )
    logging.info("===== Milvus 检索结果 =====")
    for i, r in enumerate(mv_only, 1):
        logging.info(f"Milvus Top{i} 相似度={r.get('_vector_score')} 章节={r.get('chapter')} 角色={r.get('user_role')} 页={r.get('pdf_page')}")

    # 打印融合后的混合检索结果
    logging.info("===== 混合检索（融合）结果 =====")
    for i, r in enumerate(results, 1):
        logging.info(f"Top{i} 章节={r.get('chapter')} 角色={r.get('user_role')} 页={r.get('pdf_page')}\n内容: {r.get('content','')[:200]}...")